{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 🧠 Useful Util - Convert PyTorch model's parameters and tensors to half-precision floating-point format (`float16`). \n",
    "\n",
    "This is done to save memory and improve the speed of computations at the cost of a slight loss in numerical precision."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert PyTorch model's parameters and tensors\n",
    "# to half-precision floating-point format (`float16`)\n",
    "\n",
    "from transformers.utils import logging\n",
    "\n",
    "logger = logging.get_logger(\"transformers\")\n",
    "\n",
    "\n",
    "def model_to_half(model, cast_model=True):\n",
    "    if cast_model:\n",
    "        model.half()\n",
    "\n",
    "    for n, m in model.named_modules():\n",
    "        if m.__class__.__name__ == 'QuantLinear':\n",
    "            logger.debug(f'Converting to half {n}.')\n",
    "            m.scales = m.scales.half()\n",
    "            m.bias = m.bias.half() if (m.bias is not None) else None\n",
    "    logger.info('Converted as Half.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "\n",
    "📌 First this script sets up a logger named \"transformers\", enabling you to log messages at different severity levels (debug, info, warning, etc.) in the course of the script's execution.\n",
    "\n",
    "📌 The `half()` method in PyTorch is specifically designed to convert a model's parameters and stored tensors from `float32` to `float16`. When called on a PyTorch model using `model.half()`, it recursively goes through all modules and converts all floating-point parameters and buffers to half precision.\n",
    "\n",
    "📌 The `m.scales` refers to the scale factors associated with the quantization process in these QuantLinear layers. Each scale factor is a floating-point number that is used to convert the quantized values back to their original floating-point representation.\n",
    "\n",
    "By calling `m.scales.half()`, the code is converting these scale factors from their current precision (likely float32 or full-precision floating-point) to float16. \n",
    "\n",
    "📌 If `cast_model` is `True`, the entire model is converted to half-precision floating-point format using the `half()` method provided by PyTorch.\n",
    "\n",
    "This action implicitly converts all `float32` tensors to `float16`.\n",
    "\n",
    "📌 The function then iterates through all named modules in the model using `named_modules()`.\n",
    "\n",
    "```python\n",
    "for n, m in model.named_modules():\n",
    "```\n",
    "\n",
    "Here `n` is the name of the module and `m` is the module itself.\n",
    "\n",
    "📌 Within the loop, the function checks if the module's class name is `QuantLinear`.\n",
    "\n",
    "```python\n",
    "if m.__class__.__name__ == 'QuantLinear':\n",
    "```\n",
    "\n",
    "`QuantLinear` generally refers to a linear layer designed to be quantized, i.e., to have its weights and biases represented with fewer bits than standard `float32`.\n",
    "\n",
    "📌 If the module is of type `QuantLinear`, it prints a debug log and proceeds to cast its `scales` and `bias` attributes to half-precision.\n",
    "\n",
    "```python\n",
    "logger.debug(f'Converting to half {n}.')\n",
    "m.scales = m.scales.half()\n",
    "m.bias = m.bias.half() if (m.bias is not None) else None\n",
    "```\n",
    "\n",
    "The function specifically checks whether the `bias` attribute is `None` before attempting to cast it.\n",
    "\n",
    "-----------------\n",
    "\n",
    "### 👉  Underlying Concepts\n",
    "\n",
    "Half-precision floating-point arithmetic is commonly used to accelerate training and inference. Reducing the precision of model parameters can lead to significant speedup and reduction in memory usage, especially on GPUs that are optimized for `float16` arithmetic. However, it's crucial to note that this comes at the cost of numerical precision, which may or may not impact model performance significantly depending on the application.\n",
    "\n",
    "The specialized treatment for `QuantLinear` modules suggests the function is designed to work even with quantized models. Quantization is another technique to reduce memory footprint and computational requirements. It involves approximating the floating-point numbers with a fixed number of bits, usually much fewer than 32.\n",
    "\n",
    "\n",
    "----------\n",
    "\n",
    "## ⚠ One word of caution while using `half()` method of Pytorch - ⚠\n",
    "\n",
    "If you convert the entire model to fp16, there is a chance that some of the activations functions and batchnorm layers will cause the fp16 weights to underflow, i.e., become zero. So in many situation, recommended to use `autocast` which internally converts the weights to fp32 in problematic layers.\n",
    "\n",
    "📌 Autocast can be implemented using PyTorch's `torch.cuda.amp.autocast` context manager. This feature allows automatic mixed-precision training and inference by dynamically switching between `float32` and `float16`.\n",
    "\n",
    "📌 But the approach used in the given script is the explicit conversion of the model to half-precision floating-point format (`float16`) using the `model.half()` method. So this is not the `autocast` mechanism\n",
    "\n",
    "📌 The `model.half()` method is a more straightforward but potentially riskier approach compared to `autocast`. It changes the data types of all the model's parameters and stored tensors to `float16`, which can lead to numerical instability, particularly in layers like batch normalization and certain activation functions. This is because `float16` representation has limited range and precision, making it susceptible to issues like underflow and overflow.\n",
    "\n",
    "📌 The `autocast` feature, on the other hand, performs automatic mixed-precision training. It keeps most of the network in `float16` while selectively keeping certain layers in `float32` to avoid numerical instability. With `autocast`, PyTorch automatically decides when to use `float16` and when to use `float32`, making it a safer but slightly less performant option compared to full `float16` training. Furthermore, gradient scaling is often used with `autocast` to mitigate the problem of vanishing gradients during backpropagation.\n",
    "\n",
    "📌 The script provided does not deal with these potential instabilities and issues; it simply converts the model and its QuantLinear layers to `float16`. Therefore, extra caution may be needed when using this script, particularly for models that include layers known to be problematic with `float16`, such as batch normalization."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 🧠🚫 While working with Language Models in PyTorch - If you convert the entire model to `fp16`, there is a chance that some of the activations functions and batchnorm layers will cause the `fp16` weights to underflow, i.e., become zero. \n",
    "\n",
    "🧠 So in many situation, recommended to use `autocast` which internally converts the weights to fp32 in problematic layers. 🔥 \n",
    "\n",
    "📌 Autocast can be implemented using PyTorch's `torch.cuda.amp.autocast` context manager. This feature allows automatic mixed-precision training and inference by dynamically switching between `float32` and `float16`.\n",
    "\n",
    "📌 Autocast can be implemented using PyTorch's `torch.cuda.amp.autocast` context manager. This feature allows automatic mixed-precision training and inference by dynamically switching between `float32` and `float16`.\n",
    "\n",
    "📌 To implement autocast, you need to wrap the forward pass of your model and the loss computation inside the `with autocast():` block. \n",
    "\n",
    "```python\n",
    "from torch.cuda.amp import autocast\n",
    "\n",
    "def train_one_epoch(model, optimizer, data_loader):\n",
    "    model.train()\n",
    "    for batch in data_loader:\n",
    "        inputs, labels = batch\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        with autocast():\n",
    "            outputs = model(inputs)\n",
    "            loss = loss_fn(outputs, labels)\n",
    "\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "```\n",
    "\n",
    "Here, the forward pass `outputs = model(inputs)` and the loss computation `loss = loss_fn(outputs, labels)` are wrapped inside the `with autocast():` block. PyTorch automatically casts the relevant tensors to `float16` or `float32` as needed.\n",
    "\n",
    "📌 If you're using autocast, it is also recommended to use gradient scaling to prevent underflow during the backward pass. PyTorch provides the `GradScaler` class for this purpose. Here's a modified version of the above code with gradient scaling:\n",
    "\n",
    "```python\n",
    "from torch.cuda.amp import autocast, GradScaler\n",
    "\n",
    "scaler = GradScaler()\n",
    "\n",
    "def train_one_epoch(model, optimizer, data_loader):\n",
    "    model.train()\n",
    "    for batch in data_loader:\n",
    "        inputs, labels = batch\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        with autocast():\n",
    "            outputs = model(inputs)\n",
    "            loss = loss_fn(outputs, labels)\n",
    "\n",
    "        scaler.scale(loss).backward()\n",
    "        scaler.step(optimizer)\n",
    "        scaler.update()\n",
    "```\n",
    "\n",
    "### Underlying Concepts\n",
    "\n",
    "📌 The `autocast` context manager dynamically chooses the data type (`float16` or `float32`) for each operation in the forward pass to maximize the performance while minimizing the loss of precision. Layers that are known to be numerically unstable in `float16` (like batch normalization) will automatically run in `float32`, while other operations that can safely use `float16` will do so.\n",
    "\n",
    "📌 The `GradScaler` scales the gradients before performing the optimizer step. This prevents underflows, which are numbers too small to be represented in the `float16` format, from occurring during the backpropagation. After the update, `GradScaler` adjusts the scaling factor, preparing it for the next iteration. This method is often necessary to stabilize the training when using mixed-precision.\n",
    "\n",
    "By using both autocast and gradient scaling, you can often achieve the performance benefits of `float16` training while avoiding the numerical instabilities, making it a robust and generally safer approach."
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
